package com.prime;

import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;

public class RightPrimeChecker implements CanCheckIfPrime {

    private static final WrongIfPrimeChecker delegate = new WrongIfPrimeChecker();
    private static final ConcurrentHashMap<Integer, CountDownLatch> executingNumbers = new ConcurrentHashMap<>();

    /**
     * Thread safe!
     */
    @Override
    public boolean isPrime(int number) {
        CountDownLatch lastLatch;
        CountDownLatch latch = new CountDownLatch(1);
        if ((lastLatch = executingNumbers.putIfAbsent(number, latch)) == null) {
            boolean result = delegate.isPrime(number);
            latch.countDown();
            executingNumbers.remove(number);
            return result;
        } else {
            try {
                lastLatch.await();
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
        }
        return delegate.isPrime(number);
    }
}